from torch.utils.data import DataLoader
import torchvision
from torch.utils.tensorboard import SummaryWriter

# 准备测试数据集
test_data = torchvision.datasets.CIFAR10(root="week1/PyTorch学习/dataset_CIFAR10",train = False,transform=torchvision.transforms,download=True) #测试集

test_loader = DataLoader(dataset=test_data,batch_size=4,shuffle=True,num_workers=0,drop_last=False)
# dataset = 数据集
# batch_size = 每次取出多少个数据,img_tensor[0:4]表示取出前4个数据
# shuffle = 是否打乱数据
# num_workers = 0表示不使用多线程
# drop_last = 是否丢弃最后一个batch，如果数据集不能被batch_size整除，则最后一个batch可能会小于batch_size，设置为True则丢弃这个batch
img,target = test_data[0]  # 获取第一个数据和标签
print(img.shape)  # 输出图片的形状
print(target)  # 输出标签

writer = SummaryWriter("week1/PyTorch学习/4-DataLoader的使用/logs_dataloader")  # 创建一个SummaryWriter对象，用于记录日志

# 64 + 64 + 64 + 64 + ...
# 每个data是一个64
# 输出64张，step+1，然后进行第二个64
step = 0
for data in test_loader:
    imgs,targets = data
    writer.add_image("test_data", imgs,step) 
    step = step + 1
    #print(imgs.shape)  # 输出每个batch的形状
    #print(targets)  # 输出每个batch的标签

#如果想输出两份，加一个前面的循环
#for i in range(2):
writer.close()  



